import numpy as np
import cv2

from ascend_yolo_model import AscendYOLOModel


def masks_to_contours(masks, labels):
    contours_list = []

    for i, mask in enumerate(masks.xy):  # 使用 masks.xy 获取掩码数据
        mask_np = np.array(mask)  # 将掩码转换为 NumPy 数组
        mask_np = (mask_np > 0).astype(np.uint8) * 255  # 二值化掩码并标记为255（白色）
        contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 查找轮廓

        # 提取轮廓点及其标签
        for contour in contours:
            contour_points = contour.squeeze().tolist()  # 转换为列表
            contours_list.append({
                "label": labels[i],  # 赋予当前掩码的标签
                "contour": contour_points
            })

    return contours_list


if __name__ == '__main__':
    model_path = '/app/dependences/Yolov8_for_PyTorch/model/yolov8n-seg.pt'
    om_model_path = '/app/dependences/Yolov8_for_PyTorch/model/yolov8n-seg_bs1.om'
    device_id = 0
    image_path = '/app/dependences/Yolov8_for_PyTorch/person.jpg'
    model = AscendYOLOModel(model_path=model_path,
                            task='segment',
                            om_model_path=image_path,
                            device_id=device_id)
    results = model.predict(image_path, conf=0.7, iou=0.3)
    boxes = results[0].boxes.xyxy.cpu().numpy()  # 检测框，格式为(x1, y1, x2, y2)
    scores = results[0].boxes.conf.cpu().numpy()  # 置信度
    classes = results[0].boxes.cls.cpu().numpy()  # 类别索引
    labels = results[0].names  # 类别标签
    masks = results[0].masks  # Masks 对象
    result_list = []

    # 创建一个图像来绘制结果
    img = cv2.imread(image_path)  # 读取原始图像

    for box, score, cls, mask in zip(boxes, scores, classes, masks):
        label = labels[int(cls)]
        confidence = score
        # 转换为 (height, width, channels) 格式
        mask_np = mask.cpu().data.numpy().transpose(1, 2, 0)
        mask_np = (mask_np * 255).astype(np.uint8)  # * 255   二值化掩码并标记为255（白色）
        contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 查找轮廓
        # 提取轮廓点及其标签
        if contours:
            # 找到点位最多的轮廓
            best_contour = max(contours, key=cv2.contourArea)
            # 简化轮廓，把一些相似的点位剔除。
            epsilon = 0.004 * cv2.arcLength(best_contour, True)  # 根据需要调整 epsilon 的值 True 参数表示轮廓是否是闭合的
            approx_contour = cv2.approxPolyDP(best_contour, epsilon, True)  # True 参数表示轮廓是否是闭合的
            contour_points = approx_contour.squeeze().tolist()  # 转换为列表
            result_list.append({
                "label": label,  # 赋予当前掩码的标签
                "confidence": confidence,
                "contour": contour_points
            })
            # 绘制轮廓
            cv2.drawContours(img, [approx_contour], -1, (0, 255, 255), 2)  # 绘制黄色轮廓

    # 保存带有轮廓的图像
    cv2.imwrite('result_with_contours.jpg', img)
    cv2.destroyAllWindows()

    # 打印结果
    for i in result_list:
        print(f"Label: {i.get('label')}, Confidence: {i.get('confidence')}, Contour: {i.get('contour')}")
